import logging

import torch
import torch.nn.functional as F

from trainer.base_trainer import BaseTrainer

logger = logging.getLogger("match")

class MaskLearnedTrainer(BaseTrainer):
    cnt=0
    sample=0

    # log_freq = 100 # print iter info in log with on_iteration_start
    tb_img_freq = 20 # frequency to call tensorboard log_image_grid, 0 means not use log_image_grid

    record_patch = False

    def __init__(self, model, optimizer, loss_fn, logger, config):
        super(MaskLearnedTrainer, self).__init__(
            model, optimizer, loss_fn, logger, config
        )

        self.logger = logger

        self.optim_fts = optimizer["Fts"]

        self.FtsA = model["FtsA"]
        self.FtsB = model["FtsB"] if "FtsB" in model else self.FtsA
        self.FtsM = model["FtsM"]

        self.Lmatch = loss_fn["Mask"][0]

        if not "frozen" in dir(self.opts):
            self.frozen = False
        else:
            self.frozen = self.opts.frozen
        if not "feat" in dir(self.opts):
            self.feat = "desc"
        else:
            self.feat = self.opts.feat

        self.mask_max = 1.0
        self.mask_min = 0.749019607843137
        if "max" in dir(self.opts):
            self.mask_max = self.opts.max
            self.mask_min = self.opts.min
        
        if  "record_patch" in dir(self.opts):
            self.record_patch = self.opts.record_patch


    def train(self, engine, batch):
        self.FtsA.train()
        self.FtsB.train()
        self.FtsM.train()
        
        log_str = ""
        curr_step = self.logger.counters["train"]

        (sar_patchs, opt_patchs, mask), (labels, shift) = self._prepare_batch(batch)

        self._zero_grad()

        if self.frozen:
            with torch.no_grad():
                descA, norm_descA, fmA = self.FtsA(sar_patchs)
        else:
            descA, norm_descA, fmA = self.FtsA(sar_patchs)

        featA = self.get_feat((descA, norm_descA, fmA))

        predict = self.FtsM(featA)

        mask = (mask-self.mask_min) / (self.mask_max-self.mask_min)
        loss = self._calculate_loss(predict, mask)

        self.logger.add_scalar('train/loss', loss.item(), curr_step)

        log_str += "Lmask: {:.5f} \t".format(loss.item())

        loss.backward()
        self.optim_fts.step()

        self.log_str = log_str

        return None, None


    def evaluate(self, engine, batch):
        self.FtsA.eval()
        self.FtsB.eval()
        self.FtsM.eval()

        with torch.no_grad():

            (sar_patchs, opt_patchs, mask), (labels, shift) = self._prepare_batch(batch)

            descA, norm_descA, fmA = self.FtsA(sar_patchs)
            descB, norm_descB, fmB = self.FtsB(opt_patchs)

            featA = self.get_feat((descA, norm_descA, fmA))
            featB = self.get_feat((descB, norm_descB, fmB))

            predict = self.FtsM(featA)

            loss = self._calculate_loss(predict, mask)

            h_a, h_b = norm_descA.size(2), norm_descB.size(2)
            if h_a >= h_b:
                s_p = norm_descA
                t_p = norm_descB
            else:
                s_p = norm_descB
                t_p = norm_descA

            predict = predict * (self.mask_max-self.mask_min) + self.mask_min

        corr_map = self.mtc_fn(s_p, t_p, predict)

        l2, loc = self._l2_dis((corr_map).detach(), shift, ret_pred=True)

        sampler_dict = {
            "scalar-scalar" : {
                "l2=dis" : l2, 
                "loc" : loc
            }, 
            "sar=mft-fmap_chs" : norm_descA, 
            "sar=mask-img" : (predict-self.mask_min) / (self.mask_max-self.mask_min), 
            "corr=map-fmap_chs" : corr_map,  
            "gt=mask-img" : (mask-self.mask_min) / (self.mask_max-self.mask_min),  
        }
        if self.record_patch:
            sampler_dict.update({
                "sar=img-img" : sar_patchs, 
                "opt=img-img" : opt_patchs, 
                "gt=map-img" : labels, 
            })
    
        self.logger.info(num=engine.state.iteration, **sampler_dict)

        return l2, loss


    def _calculate_loss(self, out, gt_map, *, device=""):
        if device == "":
            device = self.device

        loss = self.Lmatch(out, gt_map)
        loss = loss.to(device)

        return loss


    def get_feat(self, tensors):
        ret = None
        if self.feat == "desc":
            ret = tensors[0]
        elif self.feat == "norm_desc":
            ret = tensors[1]
        elif self.feat == "fm":
            ret = tensors[2]

        return ret







